import unittest
import sys
import json
from unittest.mock import patch, MagicMock, mock_open
import argparse

from engine_server.config.base import ServerConfig
class TestVLLMConfig(unittest.TestCase):
    
    def setUp(self):
        # Save original modules for restoration in tearDown
        self.original_modules = {}
        
        # Mock vllm-related modules
        self.mock_vllm = MagicMock()
        self.mock_vllm.utils = MagicMock()
        self.mock_vllm.utils.FlexibleArgumentParser = argparse.ArgumentParser
        
        # Mock entrypoints structure
        self.mock_entrypoints = MagicMock()
        self.mock_openai = MagicMock()
        self.mock_cli_args = MagicMock()
        self.mock_cli_args.make_arg_parser = MagicMock()
        self.mock_cli_args.validate_parsed_serve_args = MagicMock()
        
        # Save original module references and replace with mocks
        for mod_name in ['vllm', 'vllm.utils', 'vllm.entrypoints', 'vllm.entrypoints.openai', 
                        'vllm.entrypoints.openai.cli_args', 'engine_server.utils.logger']:
            if mod_name in sys.modules:
                self.original_modules[mod_name] = sys.modules[mod_name]
        
        # Set up mock modules
        sys.modules['vllm'] = self.mock_vllm
        sys.modules['vllm.utils'] = self.mock_vllm.utils
        sys.modules['vllm.entrypoints'] = self.mock_entrypoints
        sys.modules['vllm.entrypoints.openai'] = self.mock_openai
        sys.modules['vllm.entrypoints.openai.cli_args'] = self.mock_cli_args
        
        # Mock logger
        self.mock_logger = MagicMock()
        self.mock_logger.run_log = MagicMock()
        sys.modules['engine_server.utils.logger'] = self.mock_logger
        
        # Import BaseConfig to mock its methods
        if 'engine_server.config.base' in sys.modules:
            self.original_modules['engine_server.config.base'] = sys.modules['engine_server.config.base']
        
        # Import VLLMConfig and _add_argument_to_list
        from engine_server.config.vllm import VLLMConfig, _add_argument_to_list
        self.VLLMConfig = VLLMConfig
        self._add_argument_to_list = _add_argument_to_list
        
        # Set up test environment
        self.server_config = ServerConfig(
            server_host="localhost",
            server_port=9001,
            engine_type="vllm",
            config_path=None,
            dp_rank=0
        )
        self.kv_config = {
            "model": "test-model",
            "max_tokens": 1024
        }
    
    def tearDown(self):
        # Restore original modules to avoid impacting other tests
        for mod_name, original_mod in self.original_modules.items():
            sys.modules[mod_name] = original_mod
        
        # Clean up mock modules
        for mod_name in ['vllm', 'vllm.utils', 'vllm.entrypoints', 'vllm.entrypoints.openai', 
                        'vllm.entrypoints.openai.cli_args', 'engine_server.utils.logger',
                        'engine_server.config.vllm']:
            if mod_name not in self.original_modules and mod_name in sys.modules:
                del sys.modules[mod_name]
    
    def test_vllm_config_default_initialization(self):
        """Test default initialization of VLLMConfig"""
        vllm_config = self.VLLMConfig(server_config=self.server_config)
        
        # Verify basic attributes
        self.assertEqual(vllm_config.server_config, self.server_config)
        self.assertEqual(vllm_config.args, None)
        
        # Verify default mapping dictionary
        expected_mapping = {
            'dp_rank': 'data_parallel_rank',
        }
        self.assertEqual(vllm_config.mapping, expected_mapping)
        
        # Verify default kv_config is empty dictionary
        self.assertEqual(vllm_config.kv_config, {})
        
        # Verify default checkers dictionary
        self.assertIn('dp_rank', vllm_config.checkers)
        self.assertTrue(callable(vllm_config.checkers['dp_rank']))
    
    def test_vllm_config_custom_initialization(self):
        """Test custom initialization of VLLMConfig"""
        # Create custom parameters
        custom_args = argparse.Namespace(model="custom-model")
        custom_mapping = {
            'custom_key': 'custom_value',
            'dp_rank': 'data_parallel_rank'
        }
        
        # Create custom checkers
        def custom_checker(value):
            return value is not None
        
        custom_checkers = {
            'custom_key': custom_checker,
            'dp_rank': lambda x: x > 1
        }
        
        # Initialize VLLMConfig
        vllm_config = self.VLLMConfig(
            server_config=self.server_config,
            kv_config=self.kv_config,
            args=custom_args,
            mapping=custom_mapping,
            checkers=custom_checkers
        )
        
        # Verify custom attributes
        self.assertEqual(vllm_config.server_config, self.server_config)
        self.assertEqual(vllm_config.kv_config, self.kv_config)
        self.assertEqual(vllm_config.args, custom_args)
        self.assertEqual(vllm_config.mapping, custom_mapping)
        self.assertEqual(vllm_config.checkers, custom_checkers)
        
    def test_vllm_config_empty_mapping_initialization(self):
        """Test initialization with empty mapping"""
        # Initialize VLLMConfig with empty mapping
        vllm_config = self.VLLMConfig(
            server_config=self.server_config,
            mapping={}
        )
        
        # Verify empty mapping
        self.assertEqual(vllm_config.mapping, {})
    
    def test_get_param_list_basic(self):
        """Test basic functionality of _get_param_list method"""
        # Set dp_rank in server_config for testing (0 should not be included)
        self.server_config.dp_rank = 0
        
        vllm_config = self.VLLMConfig(
            server_config=self.server_config,
            kv_config=self.kv_config
        )
        
        # Get parameter list
        param_list = vllm_config._get_param_list()
        
        # Verify dp_rank=0 is NOT included in parameter list due to checker
        self.assertNotIn('--data-parallel-rank', param_list)
        
        # Verify parameter list contains attributes from kv_config
        self.assertIn('--model', param_list)
        model_index = param_list.index('--model')
        self.assertEqual(param_list[model_index + 1], 'test-model')
        
        self.assertIn('--max-tokens', param_list)
        max_tokens_index = param_list.index('--max-tokens')
        self.assertEqual(param_list[max_tokens_index + 1], '1024')
        
    def test_get_param_list_with_multiple_mappings(self):
        """Test _get_param_list with multiple mappings"""
        # Create custom server config with multiple mapped attributes
        custom_server_config = ServerConfig(
            dp_rank=2,  # dp_rank=2 should be included due to checker
            engine_type="vllm"
        )
        
        # Create custom mapping
        custom_mapping = {
            'dp_rank': 'data_parallel_rank',
            'engine_type': 'engine-type'
        }
        
        # Create custom checkers to handle engine_type
        custom_checkers = {
            'dp_rank': lambda x: x > 0,  # Use default dp_rank checker
            'engine_type': lambda x: x == "vllm"  # Only include if engine_type is vllm
        }
        
        # Initialize VLLMConfig
        vllm_config = self.VLLMConfig(
            server_config=custom_server_config,
            mapping=custom_mapping,
            checkers=custom_checkers
        )
        
        # Get parameter list
        param_list = vllm_config._get_param_list()
        
        # Verify both mappings are included (dp_rank=2 passes checker)
        self.assertIn('--data-parallel-rank', param_list)
        dp_rank_index = param_list.index('--data-parallel-rank')
        self.assertEqual(param_list[dp_rank_index + 1], '2')
        
        self.assertIn('--engine-type', param_list)
        engine_type_index = param_list.index('--engine-type')
        self.assertEqual(param_list[engine_type_index + 1], 'vllm')
        
    def test_get_param_list_missing_attribute(self):
        """Test _get_param_list when server_config is missing a mapped attribute"""
        # Create VLLMConfig with mapping for non-existent attribute
        vllm_config = self.VLLMConfig(
            server_config=self.server_config,
            mapping={'non_existent_attr': 'mapped_name'}
        )
        
        # Get parameter list
        param_list = vllm_config._get_param_list()
        
        # Verify missing attribute doesn't appear in the list
        self.assertNotIn('--mapped-name', param_list)
    
    def test_get_param_list_with_different_types(self):
        """Test _get_param_list method handles values of different types"""
        # Create configs with various value types
        custom_server_config = ServerConfig(
            dp_rank=1  # Test mapping with int value
        )
        custom_kv_config = {
            "boolean_true": True,
            "boolean_false": False,
            "string_value": "test-string",
            "int_value": 42,
            "float_value": 3.14,
            "list_value": ["item1", "item2"]
        }
        
        vllm_config = self.VLLMConfig(
            server_config=custom_server_config,
            kv_config=custom_kv_config
        )
        
        # Get parameter list
        param_list = vllm_config._get_param_list()
        
        # Verify parameter list contains mapped dp_rank
        self.assertIn('--data-parallel-rank', param_list)
        dp_rank_index = param_list.index('--data-parallel-rank')
        self.assertEqual(param_list[dp_rank_index + 1], '1')
        
        # Verify boolean true is added as flag only
        self.assertIn('--boolean-true', param_list)
        self.assertNotIn('True', param_list)  # Boolean true should just be the flag
        
        # Verify boolean false is not added
        self.assertNotIn('--boolean-false', param_list)
        
        # Verify string value
        self.assertIn('--string-value', param_list)
        string_index = param_list.index('--string-value')
        self.assertEqual(param_list[string_index + 1], 'test-string')
        
        # Verify int value
        self.assertIn('--int-value', param_list)
        int_index = param_list.index('--int-value')
        self.assertEqual(param_list[int_index + 1], '42')
        
        # Verify float value is converted to string
        self.assertIn('--float-value', param_list)
        float_index = param_list.index('--float-value')
        self.assertEqual(param_list[float_index + 1], '3.14')
        
        # Verify list values
        self.assertIn('--list-value', param_list)
        list_index = param_list.index('--list-value')
        self.assertEqual(param_list[list_index + 1], 'item1')
        self.assertEqual(param_list[list_index + 2], 'item2')
        
    def test_get_param_list_with_empty_values(self):
        """Test _get_param_list method handles empty values"""
        # Create configs with empty values
        custom_kv_config = {
            "empty_string": "",
            "empty_list": [],
            "none_value": None
        }
        
        vllm_config = self.VLLMConfig(
            server_config=self.server_config,
            kv_config=custom_kv_config
        )
        
        # Get parameter list
        param_list = vllm_config._get_param_list()
        
        # Verify empty list is not added
        self.assertNotIn('--empty-list', param_list)
        
        # Verify empty string and None values are still added
        self.assertIn('--empty-string', param_list)
        empty_str_index = param_list.index('--empty-string')
        self.assertEqual(param_list[empty_str_index + 1], '')
        
        self.assertIn('--none-value', param_list)
        none_index = param_list.index('--none-value')
        self.assertEqual(param_list[none_index + 1], 'None')
    
    def test_add_argument_to_list_function(self):
        """Test _add_argument_to_list helper function"""
        # Test boolean True
        arg_list = []
        self._add_argument_to_list(arg_list, 'test-boolean', True)
        self.assertEqual(arg_list, ['--test-boolean'])
        
        # Test boolean False
        arg_list = []
        self._add_argument_to_list(arg_list, 'test-boolean', False)
        self.assertEqual(arg_list, [])
        
        # Test string value
        arg_list = []
        self._add_argument_to_list(arg_list, 'test-string', 'value')
        self.assertEqual(arg_list, ['--test-string', 'value'])
        
        # Test integer value
        arg_list = []
        self._add_argument_to_list(arg_list, 'test-int', 42)
        self.assertEqual(arg_list, ['--test-int', '42'])
        
        # Test list value
        arg_list = []
        self._add_argument_to_list(arg_list, 'test-list', ['item1', 'item2'])
        self.assertEqual(arg_list, ['--test-list', 'item1', 'item2'])
        
        # Test empty list
        arg_list = []
        self._add_argument_to_list(arg_list, 'test-empty-list', [])
        self.assertEqual(arg_list, [])
        
        # Test float value
        arg_list = []
        self._add_argument_to_list(arg_list, 'test-float', 3.14)
        self.assertEqual(arg_list, ['--test-float', '3.14'])
        
        # Test None value
        arg_list = []
        self._add_argument_to_list(arg_list, 'test-none', None)
        self.assertEqual(arg_list, ['--test-none', 'None'])
        
        # Test empty string
        arg_list = []
        self._add_argument_to_list(arg_list, 'test-empty-string', '')
        self.assertEqual(arg_list, ['--test-empty-string', ''])
    
    @patch('engine_server.config.base.BaseConfig.convert')
    def test_convert_method(self, mock_base_convert):
        """Test convert method"""
        # Create VLLMConfig instance
        vllm_config = self.VLLMConfig(server_config=self.server_config, kv_config=self.kv_config)
        
        # Call convert method
        vllm_config.convert()
        
        # Assert that BaseConfig.convert is called
        mock_base_convert.assert_called_once()
        
    @patch('engine_server.config.base.BaseConfig.convert')
    @patch('engine_server.config.vllm.sys')
    @patch('engine_server.config.vllm.FlexibleArgumentParser')
    @patch('engine_server.config.vllm.make_arg_parser')
    @patch('engine_server.config.vllm.run_log')
    def test_convert_method_without_config_file(self, mock_run_log, mock_make_arg_parser, mock_flexible_parser, mock_sys, mock_base_convert):
        """Test convert method when no config file is specified"""
        # Create mock parser and parameters
        mock_parser_instance = MagicMock()
        mock_flexible_parser.return_value = mock_parser_instance
        mock_make_arg_parser.return_value = mock_parser_instance
        
        # Create expected parsing result
        expected_args = argparse.Namespace(model="parsed-model")
        mock_parser_instance.parse_args.return_value = expected_args
        
        # Set dp_rank but no config_path (dp_rank=0 will be excluded by checker)
        self.server_config.dp_rank = 0
        self.server_config.config_path = None
        
        # Initialize VLLMConfig
        vllm_config = self.VLLMConfig(
            server_config=self.server_config,
            kv_config=self.kv_config
        )
        
        # Call convert method
        vllm_config.convert()
        
        # Verify parent convert method is called
        mock_base_convert.assert_called_once()
        
        # Verify args are set correctly
        self.assertEqual(vllm_config.args, expected_args)
    
    def test_should_include_param_method(self):
        """Test _should_include_param method with different checkers"""
        # Initialize VLLMConfig
        vllm_config = self.VLLMConfig(server_config=self.server_config)
        
        # Test with dp_rank checker (should return False for 0)
        self.assertFalse(vllm_config._should_include_param('dp_rank', 0))
        # Test with dp_rank checker (should return True for positive values)
        self.assertTrue(vllm_config._should_include_param('dp_rank', 1))
        
        # Test with default checker (should return True for any value)
        self.assertTrue(vllm_config._should_include_param('non_existent_key', 'any_value'))
        
        # Test with custom checkers
        def custom_checker(value):
            return isinstance(value, str) and len(value) > 5
        
        vllm_config.checkers['custom_key'] = custom_checker
        
        # Test custom checker
        self.assertTrue(vllm_config._should_include_param('custom_key', 'long_string'))
        self.assertFalse(vllm_config._should_include_param('custom_key', 'short'))
        self.assertFalse(vllm_config._should_include_param('custom_key', 123))
    
    @patch('engine_server.config.base.BaseConfig.convert')
    def test_convert_method_error_handling(self, mock_base_convert):
        """Test error handling in convert method"""
        # Create VLLMConfig instance
        vllm_config = self.VLLMConfig(server_config=self.server_config, kv_config=self.kv_config)
        
        # Mock BaseConfig.convert to raise an exception
        mock_base_convert.side_effect = ValueError("Test error")
        
        # Test that the exception is propagated
        with self.assertRaises(ValueError) as context:
            vllm_config.convert()
        
        self.assertEqual(str(context.exception), "Test error")
        
    @patch('engine_server.config.base.BaseConfig.convert')
    def test_convert_method_base_convert_exception(self, mock_base_convert):
        """Test convert method when BaseConfig.convert raises an exception"""
        # Simulate BaseConfig.convert raising exception
        mock_base_convert.side_effect = ValueError("Base config error")
        
        # Initialize VLLMConfig
        vllm_config = self.VLLMConfig(server_config=self.server_config)
        
        # Verify exception is propagated correctly
        with self.assertRaises(ValueError) as context:
            vllm_config.convert()
        
        self.assertEqual(str(context.exception), "Base config error")
    
    def test_get_args_method(self):
        """Test get_args method"""
        # Create test arguments
        test_args = argparse.Namespace(model="test-model", max_tokens=1024)
        
        # Initialize VLLMConfig and set args
        vllm_config = self.VLLMConfig(server_config=self.server_config)
        vllm_config.args = test_args
        
        # Verify get_args returns correct args
        result = vllm_config.get_args()
        self.assertEqual(result, test_args)
        
        # Verify same object is returned
        self.assertIs(result, test_args)
        
        # Test case where args is None
        vllm_config = self.VLLMConfig(server_config=self.server_config)
        vllm_config.args = None
        self.assertIsNone(vllm_config.get_args())
    
    def test_get_server_config_method(self):
        """Test get_server_config method"""
        # Initialize VLLMConfig
        vllm_config = self.VLLMConfig(server_config=self.server_config)
        
        # Verify get_server_config returns correct server_config
        result = vllm_config.get_server_config()
        self.assertEqual(result, self.server_config)
        
        # Verify same object is returned
        self.assertIs(result, self.server_config)
        
        # Test with different server_config
        different_server_config = ServerConfig(
            server_host="different-host",
            server_port=9002,
            engine_type="vllm",
            config_path=None,
            dp_rank=1
        )
        vllm_config = self.VLLMConfig(server_config=different_server_config)
        self.assertEqual(vllm_config.get_server_config(), different_server_config)
    
    @patch('engine_server.config.base.BaseConfig.validate')
    @patch('engine_server.config.vllm.validate_parsed_serve_args')
    def test_validate_method_success(self, mock_validate_parsed_serve_args, mock_base_validate):
        """Test validate method behavior under normal conditions"""
        # Initialize VLLMConfig
        vllm_config = self.VLLMConfig(server_config=self.server_config)
        vllm_config.args = argparse.Namespace(model="test-model")
        
        # Call validate method
        vllm_config.validate()
        
        # Verify parent validate method is called
        mock_base_validate.assert_called_once()
        
        # Verify validate_parsed_serve_args is called
        mock_validate_parsed_serve_args.assert_called_once_with(vllm_config.args)
        
    @patch('engine_server.config.base.BaseConfig.validate')
    def test_validate_method_no_args(self, mock_base_validate):
        """Test validate method when args is None"""
        # Initialize VLLMConfig with no args
        vllm_config = self.VLLMConfig(server_config=self.server_config)
        vllm_config.args = None
        
        # Call validate method should still work (validate_parsed_serve_args might handle None)
        vllm_config.validate()
        
        # Verify parent validate method is called
        mock_base_validate.assert_called_once()
    
    @patch('engine_server.config.base.BaseConfig.validate')
    @patch('engine_server.config.vllm.validate_parsed_serve_args')
    def test_validate_method_with_invalid_args(self, mock_validate_parsed_serve_args, mock_base_validate):
        """Test validate method behavior with invalid arguments"""
        # Simulate validate_parsed_serve_args throwing exception
        mock_validate_parsed_serve_args.side_effect = ValueError("Invalid arguments")
        
        # Initialize VLLMConfig
        vllm_config = self.VLLMConfig(server_config=self.server_config)
        vllm_config.args = argparse.Namespace(model="invalid-model")
        
        # Verify exception is propagated correctly
        with self.assertRaises(ValueError) as context:
            vllm_config.validate()
        
        self.assertEqual(str(context.exception), "Invalid arguments")
        mock_base_validate.assert_called_once()
        mock_validate_parsed_serve_args.assert_called_once_with(vllm_config.args)
    
    @patch('engine_server.config.base.BaseConfig.validate')
    @patch('engine_server.config.vllm.validate_parsed_serve_args')
    def test_validate_method_with_invalid_port(self, mock_validate_parsed_serve_args, mock_base_validate):
        """Test validate method behavior with invalid port"""
        # Simulate BaseConfig.validate throwing exception
        mock_base_validate.side_effect = ValueError("Invalid port")
        
        # Initialize VLLMConfig
        vllm_config = self.VLLMConfig(server_config=self.server_config)
        
        # Verify exception is thrown correctly (from BaseConfig's validate method)
        with self.assertRaises(ValueError) as context:
            vllm_config.validate()
        
        self.assertEqual(str(context.exception), "Invalid port")
        # validate_parsed_serve_args should not be called due to port validation failure
        mock_validate_parsed_serve_args.assert_not_called()
        
    def test_integration_convert_and_validate(self):
        """Test integration of convert and validate methods"""
        # Create a test scenario where convert sets up args that validate will check
        with patch('engine_server.config.base.BaseConfig.convert'), \
             patch('engine_server.config.base.BaseConfig.validate'), \
             patch('engine_server.config.vllm.sys'), \
             patch('engine_server.config.vllm.FlexibleArgumentParser'), \
             patch('engine_server.config.vllm.make_arg_parser'), \
             patch('engine_server.config.vllm.run_log'), \
             patch('engine_server.config.vllm.validate_parsed_serve_args') as mock_validate:
            
            # Create mock parser and args
            mock_parser_instance = MagicMock()
            args = argparse.Namespace(model="test-model")
            
            # Set up mocks
            mock_flexible_parser = patch('engine_server.config.vllm.FlexibleArgumentParser')
            mock_flexible_parser.return_value = mock_parser_instance
            mock_parser_instance.parse_args.return_value = args
            
            # Initialize and convert
            vllm_config = self.VLLMConfig(server_config=self.server_config)
            vllm_config.args = args  # Manually set args to simulate convert
            
            # Validate
            vllm_config.validate()
            
            # Verify validation was called with the converted args
            mock_validate.assert_called_once_with(args)


if __name__ == '__main__':
    unittest.main()